import imaplib
import email.parser
import email.message

from mailstore import MailStore, MailFolder

def ImapMessageFactory(folder,uid,ms,unread=False, part=None):
    """Produce a factory for ImapMessage instances.

    folder, uid, ms, unread and part are the same as documented in
    ImapMessage.

    Returns a no argument callable that produces ImapMessage objects
    with the given message parameters.
    
    This is primarily useful for email.parser.Parser instances.
    """
    return lambda: ImapMessage(folder,uid,ms,unread, part)

class ImapMessage(email.message.Message, object):
    """A delayed loading, IMAP capable Message object.

    A subclass of email.message.Message that get's the payload dynamically
    from an Imap4MailStore instance.
    """
    def __init__(self, folder, uid, ms, unread=False, part=None):
        super(ImapMessage, self).__init__()
        self.folder = folder
        self.uid = uid
        self.part = part
        self.unread = unread
        self._ms = ms
        self._payload_loaded = False
        print "imapmessage %s in %s part %s" % (self.uid, self.folder, self.part)

    def get_payload(self, i=None, decode=False):
        """
        An overloaded version of Message.get_payload() that
        retrieves the payload on demand from an Imap4MailStore.
        """
        if self._payload_loaded:
            return super(ImapMessage, self).get_payload(i, decode)

        body = ''

        if self.is_multipart():
            # We need to pull the MIME for all of our subparts
            # and make new ImapMessages for them, setting the
            # part numbers as appropriate.
            #
            # This may also need special handling for
            # message/rfc822

            # we need to pull the part numbers for our sub parts

            body = list()
            print "getting part count for %s" % self.part
            pnum = self._ms.message_part_count(self.folder, self.uid, self.part)

            for c in range(1, pnum + 1):
                if self.part is not None:
                    pstr = ''.join([self.part, '.', str(c)])
                else:
                    pstr = str(c)
                m = self._ms.message_head(self.folder, self.uid, pstr)
                if m is not None:
                    body.append(m)
        else:
            print "getting simple payload for %s" % self.part
            body = self._ms.message_body(self.folder, self.uid, self.part)

        self.set_payload(body)
        self._payload_loaded = True
        return super(ImapMessage, self).get_payload(i, decode)

    def is_multipart(self):
        """Is this a multipart message?

        An overloaded version of Message.is_multipart() that bases
        it's decision on header content rather than the payload
        we haven't loaded yet.
        """

        # Special case to handle 'message/delivery-status', plus it won't hurt anything.
        if self._payload_loaded:
            return super(ImapMessage, self).is_multipart()

        if not self.has_key('content-type'):
            return False

        ### FIXME: This may not work well for embedded messages,
        # May need special handling for embedded message/rfc822
        # May suffer from general wierdness.
        ### FIXME: Lowercase before comparison.
        if self['content-type'].startswith('multipart/') or \
            self['content-type'] == 'message/rfc822':
            return True
        else:
            return False

class Imap4MailStore(MailStore):
    """An IMAP4rev1 capable MailStore

    A MailStore that know how to handle IMAP4. Using methods that are not
    inherited from MailStore should be limited to other IMAP aware components.
    Ideally the only class that should use them is ImapMessage.

    This MailStore was written against Cyrus IMAP version 2.2. It MAY work with others.

    **This class heavily relies on stable UIDs.**
    Really. A lot. Across logins even. Nothing will work without them.
    On the bright side, we represent them as strings so they can be as
    long, wierd or complex as needed.

    This class requires your IMAP server to support the following operations:
        NAMESPACE
        SELECT: This must return NO on nonselectable boxes, not BAD
        STATUS MESSAGES: This should also return NO for bad boxes.
        LIST %: This should return all folders
        LIST folder %<delim>%: This must return all child folders
        UID SEARCH UNSEEN
        UID SORT (REVERSE ARRIVAL)
        UID FETCH RFC822.HEADER
        UID FETCH TEXT
        UID FETCH BODYSTRUCTURE: Chances are good that if this differs from cyrus nothing will work at all
        UID FETCH BODY[TEXT]
        UID FETCH BODY[<part>]
        UID FETCH BODY[<part>.MIME]: for practically any part
        UID FETCH BODY[<part>.HEADER]: for message/rfc822
    """
    _totals = {}
    _unread = {}
    _selected = None
    _structure = {}

    def __init__(self, host, port):
        self._delim = '.'
        if port is None:
            self._M = imaplib.IMAP4(host)
        else:
            self._M = imaplib.IMAP4(host, port)
#        self._M.debug=4

    def _clear_cache(self):
        """Clear the cache

        All internal state that this class has accumulated
        MUST be invalidated by this method.
        """

#        print "Clearing cache"
        self._totals = {}
        self._unread = {}
        self._structure = {}
        self._selected = None

    def login(self, username, password):
        """Login a user.

        This will clear the cache before imaplib.login is called, so
        use with care.
        """

#        print "Login called"
        self._clear_cache()
        self._M.login(username,password)

    def logout(self):
        """Logout a user.

        This will clear the cache before logout.
        """

#        print "Logout called"
        self._clear_cache()
        self._M.logout()

    def namespaces(self):
        """Returns a list of namespace strings.
        
        Currently all namespaces are in one list.
        This is likely to change in the future.
        """

        rlist = list()
        res, data = self._M.namespace()
        if res == 'OK':
            for i in data[0].strip('()').split(' (('):
                name, delim = i.split()
                rlist.append((name.strip('"'),delim.strip('"')))
        return rlist

    def _select(self, name):
        """Select a mailbox.

        This is safe to call multiple times, as it will not make
        calls for the currently selected mailbox.

        returns True on success, False otherwise.
        """

        if self._selected == name:
#            print "%s already selected" % name
            return True
#        print "selecting %s" % name
        res, data = self._M.select(name)
        if res == 'OK':
#            print "select returned: %s" % data
            self._totals[name] = int(data[0])
            self._selected = name
            return True
#        else:
#            print "couldn't select %s" % name
        return False

    def open(self, name):
        """Open a mail folder.

        Open will select the mailbox and return a MailFolder instance.
        returns a MailFolder instance regardless of success.
        """

        self._select(name)
        return MailFolder(name, self)

    def total(self, name):
        """Get the number of messages in a mailbox.

        This does a STATUS query on the given mailbox to determine
        the total number of messages.
        Calling this multiple times is safe, as it caches the result.

        returns an int representing the total number of messages.
        Any failure will result in a total of 0.
        """

        try:
            totals = self._totals[name]
#            print "returning total %d for %s from cache" % (totals, name)
            return totals
        except KeyError:
#            print "totals for %s NOT found in cache" % name
            self._totals[name] = 0

        res, data = self._M.status(name, '(MESSAGES)')
        if res == 'OK':
            start = data[0].rfind('MESSAGES ')
            start += 9
            end = data[0].find(')', start)
            self._totals[name] = int(data[0][start:end])
        return self._totals[name]

    def unread_messages(self, name):
        """
        Issues a UID SEARCH UNSEEN to determine the uids of
        unread messages in a folder.
        It is safe to call this multiple times as the result is
        cached.
        returns a list of strings representing uids. On failure, an
        empty list is returned.
        """

        try:
            unread = self._unread[name]
#            print "returning unread list for %s from cache" % name
            return unread
        except KeyError:
#            print "unread list for %s NOT in cache" % name
            self._unread[name] = list()

        if not self._select(name):
            return self._unread[name]

        res, data = self._M.uid('SEARCH', None, 'UNSEEN')
        if res == 'OK':
            self._unread[name] = data[0].split()
        return self._unread[name]

    def unread(self, name):
        """Returns the number of unread messages in a folder.

        Actually returns the len() of the result from unread_messages().
        As such any errors will return zero.
        """

        return len(self.unread_messages(name))

    def messages(self, name, start, count):
        """
        Get the first 'count' messages starting from 'start' sorted by
        REVERSE ARRIVAL.
        This will return a list of ImapMessage objects.
        On failure it will return an empty list.
        This is not currently safe to call multiple times.
        """

        ### FIXME: This should cache the sort result. ###
        messages = list()
        ulist = self.unread_messages(name)
        if not self._select(name):
            return messages

        res, data = self._M.uid('SORT', '(REVERSE ARRIVAL)', 'US-ASCII', 'ALL')
        if res == 'OK':
            mlist = data[0].split()
            for i in mlist[start:start+count]:
                unseen = False
                if i in ulist:
                    unseen = True
                messages.append(self.message(name, i, unseen))
        return messages

    def folders(self, namespace=None):
        """ Returns a list of child mailfolders.
        """
        folders = list()
        res, data = self._M.list(namespace, '%')

        if res != 'OK':
            return None

        if data[0] is None:
            return None

        for f in data:
#            print "Parsing %s" % f
            delim = f.split()[1][1:-1]
            name = f.split()[2][1:-1]
            folders.append(MailFolder(name, self))
        return folders

    def parent(self, name):
        """
        return a MailFolder instance corresponding to the parent folder.
        This method is astoundingly dumb, merely reparsing the folder
        by the delimiter.
        if there is no parent folder, e.g. for INBOX, None is returned.
        """
        ### FIXME: This should be namespace aware. ###
        p = None
        partlist = name.rsplit(self._delim, 1)
        if len(partlist) > 1:
            p = MailFolder(partlist[0], self)
        return p

    def children(self, name):
        """
        return a list of child MailFolders on success, None otherwise.
        """
        ### FIXME: this should be namespace aware. ###
        folders = list()
        res, data = self._M.list(name, '%'.join([self._delim, '%']))

        if res != 'OK':
            return None

        if data[0] is None:
            return None

        for f in data:
#            print "Parsing %s" % f
            delim = f.split()[1][1:-1]
            name = f.split()[-1][1:-1]
            folders.append(MailFolder(name, self))
        return folders

    def message(self, folder, uid, unseen=False):
        """
        Create an ImapMessage for a given uid and folder.
        This ImapMessage will be loaded with full headers.
        Returns None on error.
        """
        parser = email.parser.Parser(ImapMessageFactory(folder,uid,self, unseen))
        print "getting headers for %s %s" % (folder, uid)
        if not self._select(folder):
            return None

        res, data = self._M.uid('FETCH', uid, 'RFC822.HEADER')
        if res == 'OK':
                return parser.parsestr(data[0][-1], True)

    def _find_part(self, part, start, bs):
        """
        Return the starting string index for the part in question, given a body
        structure and part number.
        part is an int.
        start is where to begin inside of bs.
        bs is a string representing the bodystructure

        if the part cannot be found, -1 is returned. If no part is passed, start
        is returned.
        """
#        print "finding part %s, start %s in %s" % (part, start, bs)

        if part is None:
#            print "returning start"
            return start

        count = 0
        depth = 0
        pos = start

        for i in range(start, len(bs)):
            if bs[i] == '"' and depth == 0:
                break
            if bs[i] == '(':
#                print "( at %s" % i
                if depth == 0:
                    pos = i + 1
#                    print "new pos %s" % i
                depth += 1
            if bs[i] == ')':
                depth -= 1
#        print ") at %s, depth %s" % (i, depth)
                if depth == 0:
                    count += 1
#                    print "new count %s" % count
                    if count == part:
#                        print "found part %s at %s" % (part, pos)
                        return pos
#        print "returning -1"
        return -1

    def message_part_count(self, folder, uid, part=None):
        """Find the number of message parts under the given part.

        part is a string in imap part number format.

        Returns the number of submessage objects for a given imap message
        and part.

        On error, 0 is returned.
        """
        
#        print "message_part_count called for part %s" % part

        bs = self._bodystructure(folder, uid)
        if bs is None:
            return 0
        depth = 0
        count = 0


        # Rather than actually parse BODYSTRUCTUREs, we'll just pull the count of
        # message parts under the one in question. We do this by counting
        # descents and ascents in the BODYSTRUCTURE. Whenever our depth
        # returns to the part in question we increment the counter. We stop
        # counting when we reach the mime subtype for our part so we don't
        # count extra mime params.

        start = 0
        ### find the part we're interested in.
        if part is not None:
            if part.rfind('.') != -1:
                part_list = part[:part.rfind('.')].split('.')
            else:
                part_list = [part]
            for p in part_list:
                start = self._find_part(int(p), start, bs)
                if start == -1:
                    return 0

        if part is None:
            part = '1'

    # We need to get the last part number from strings of the form '1' or '1.2.3'
        if part.rfind('.') != -1:
            part = int(part[part.rfind('.'):])
        else:
                part = int(part)

        for i in range(start, len(bs)):
            if bs[i] == '"' and depth == part:
                break
            if bs[i] == '(':
                depth += 1
            if bs[i] == ')':
                depth -= 1
                if depth == part:
                    count += 1
#        print 'counted %s parts of depth %s in %s' % (count, part, bs[start:])
        return count

    def _bodystructure(self, folder, uid):
        """
        Return the content of the BODYSTRUCTURE for a given message uid
        in a folder. This will return a string with the structure itself as content.
        It is safe to call this multiple times, as it caches responses.
        """
        try:
            r = self._structure[folder][uid]
            print 'returning bodystructure from cache'
            return r
        except KeyError:
            print 'getting bodystructure'
            if not self._select(folder):
                return None
            res, data = self._M.uid('FETCH', uid, 'BODYSTRUCTURE')
            if res != 'OK':
                return None
            start = data[0].find('BODYSTRUCTURE')
            start = data[0].find('(', start)
            try:
                self._structure[folder][uid] = data[0][start:-1]
            except KeyError:
                self._structure[folder] = {}
                self._structure[folder][uid] = data[0][start:-1]
        return self._structure[folder][uid]

    def message_body(self, folder, uid, part=None):
        """
        Returns a string containing the BODY[TEXT] of a message. This will be the
        full body if the message is not multipart.
        If part is none BODY[TEXT] is provided, otherwise BODY[part].
        """
        if not self._select(folder):
            return None
        if part is None:
            res, body = self._M.uid('FETCH', uid, '(BODY[TEXT])')
        else:
            res, body = self._M.uid('FETCH', uid, ''.join(['(BODY[', part, '])']))
        if res != 'OK':
            return None
        return body[0][1]

    def message_head(self, folder, uid, part):
        """Return a message object for a given part.

        The part is returned from BODY[part.MIME]
        If the part is a message/delivery-status a full BODY[part] is returned
        """
        print "fetching mime data for %s" % part
        parser = email.parser.Parser(ImapMessageFactory(folder,uid,self, unread=False, part=part))
        if not self._select(folder):
            return None
        res, head = self._M.uid('FETCH', uid, '(BODY[' + part + '.MIME])')
        if res != 'OK':
            return None
        m = parser.parsestr(head[0][1])

        # Handle oddity of the python email libs. They really want these in special 
        # format, so get the full part at once and have it parse it.
        full_message_types=[
            'message/delivery-status',
            'message/rfc822'
        ]

        if m['content-type'] in full_message_types:
            res, body = self._M.uid('FETCH', uid, '(BODY[' + part + '])')
            if res != 'OK':
                return None
            m = parser.parsestr(''.join([head[0][1], body[0][1]]))
            m._payload_loaded = True
        return m

class Imap4SSLMailStore(Imap4MailStore):
    """
    An IMAP4MailStore that connects via SSL.
    """
    def __init__(self, host, port):
        self._delim = '.'
        if port is None:
            self._M = imaplib.IMAP4_SSL(host)
        else:
            self._M = imaplib.IMAP4_SSL(host, port)
